# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/24 09:29 
@Author  : ZhangShenao 
@File    : line_delimiter_document_loader.py 
@Desc    : 按行分割的文档加载器
"""
from typing import Iterator, AsyncIterator

import aiofiles
from langchain_core.document_loaders import BaseLoader
from langchain_core.documents import Document


class LineDelimiterDocumentLoader(BaseLoader):
    """
    自定义文档加载器,按行分割加载文档
    """

    def __init__(self, file_path: str) -> None:
        """构造函数"""

        # 初始化文件路径
        self.__file_path = file_path

    def lazy_load(self) -> Iterator[Document]:
        """按行分割加载文档"""

        # 记录行号
        line_no = 1

        # 读取文件
        with open(self.__file_path, mode='r', encoding='utf-8') as f:
            # 按行加载,生成Document并返回
            for line in f:
                yield Document(page_content=line,
                               metadata={'source': self.__file_path,
                                         'line_no': line_no,
                                         }
                               )

                line_no += 1

    async def alazy_load(self) -> AsyncIterator[Document]:
        """异步实现,按行分割加载文档"""

        # 记录行号
        line_no = 1

        # 异步读取文件
        with aiofiles.open(self.__file_path, mode='r', encoding='utf-8') as f:
            # 按行加载,生成Document并返回
            async for line in f:
                yield Document(page_content=line,
                               metadata={'source': self.__file_path,
                                         'line_no': line_no,
                                         }
                               )

                line_no += 1
